import org.w3c.dom.Document;
import org.w3c.dom.Element;
import org.w3c.dom.NodeList;

import javax.xml.parsers.DocumentBuilder;
import javax.xml.parsers.DocumentBuilderFactory;
import javax.xml.transform.Transformer;
import javax.xml.transform.TransformerFactory;
import javax.xml.transform.dom.DOMSource;
import javax.xml.transform.stream.StreamResult;
import java.io.FileInputStream;
import java.util.*;

/**
 * Describes transformation rule.
 */
public class Rule
{
    /**
     * Describes condition and transformation rule for each token.
     */
    public static class Part
    {
        public Tag originalTag;

        public Tag correctTag;

        /*
         * Inherited.
         */

        @Override
        public boolean equals(Object o)
        {
            if (this == o)
                return true;
            if (o == null || getClass() != o.getClass())
                return false;

            Part part = (Part) o;

            if (correctTag != null ? !correctTag.equals(part.correctTag) : part.correctTag != null)
                return false;
            if (!originalTag.equals(part.originalTag))
                return false;

            return true;
        }

        @Override
        public int hashCode()
        {
            int result = originalTag.hashCode();
            result = 31 * result + (correctTag != null ? correctTag.hashCode() : 0);
            return result;
        }

        @Override
        public String toString()
        {
            StringBuffer buffer = new StringBuffer();
            buffer.append(originalTag.toString());
            if (correctTag != null)
            {
                buffer.append(" -> ");
                buffer.append(correctTag.toString());
            }
            return buffer.toString();
        }
    }

    /**
     * Array of parts.
     */
    public Part parts[];

    /**
     * Create empty rule
     *
     * @param length length of rule
     */
    public Rule(int length)
    {
        parts = new Part[length];
        for (int i = 0; i < length; i++)
            parts[i] = new Part();
    }

    public int Length()
    {
        return parts.length;
    }

    public static class TestResult //implements Comparable
    {
        public int total = 0;

        public int right = 0;

        public int wrong = 0;

        public TestResult()
        {
        }

        public void Add(TestResult testResult)
        {
            this.total += testResult.total;
            this.right += testResult.right;
            this.wrong += testResult.wrong;
        }

        @Override
        public String toString()
        {
            return String.format("{total: %d, right: %d, wrong: %d}", total, right, wrong);
        }
    }

    public TestResult Test(List<Token> tokens)
    {
        TestResult result = new TestResult();
        for (int i = 0; i < tokens.size() - Length() + 1; i++)
        {
            //Rule.TestResult singleResult = ApplyToIndex(tokens, i, false, null);
            Rule.TestResult singleResult = TestAtIndex(tokens, i, null);
            result.Add(singleResult);
        }
        return result;
    }

    public TestResult Apply(List<Token> tokens, Collection<Rule> newRules)
    {
        TestResult result = new TestResult();
        LinkedList<Integer> indexesToApply = new LinkedList<Integer>();
        HashSet<Integer> indexesToLearn = null;
        if (newRules != null)
            indexesToLearn = new HashSet<Integer>();
        for (int i = 0; i < tokens.size() - Length() + 1; i++)
        {
            //Rule.TestResult singleResult = ApplyToIndex(tokens, i, true, changedIndexes);
            Rule.TestResult singleResult = TestAtIndex(tokens, i, indexesToLearn);
            if (singleResult.total > 0)
                indexesToApply.add(i);
            result.Add(singleResult);
        }
        if (newRules != null)
            for (Integer tokenIndex : indexesToLearn)
                LearnAllAroundIndex(newRules, tokens, tokenIndex);
        for (int i : indexesToApply)
            ApplyAtIndex(tokens, i);
        if (newRules != null)
            for (Integer tokenIndex : indexesToLearn)
                LearnAllAroundIndex(newRules, tokens, tokenIndex);
        return result;
    }

    private TestResult TestAtIndex(List<Token> tokens, int startIndex,
                                   Collection<Integer> changedIndexes)
    {
        // check bounds
        assert 0 <= startIndex && startIndex + Length() <= tokens.size();

        // test if rule can be applied and calculate result
        TestResult result = new TestResult();
        boolean canApply = true;
        HashSet<Integer> localChangedIndexes = new HashSet<Integer>();
        for (int i = 0; i < Length(); i++)
        {
            Part part = parts[i];
            int index = startIndex + i;
            Token token = tokens.get(index);

            // current tag doesn't match - can't apply rule
            if (!token.currentTag.equals(part.originalTag))
            {
                canApply = false;
                break;
            }

            // correct tag not present in list of possible tags - rule can't be applied
            if (part.correctTag != null && !token.tagSet.contains(part.correctTag))
            {
                canApply = false;
                break;
            }

            // nothing to correct for this part
            if (part.correctTag == null)
                continue;
            else
            {
                assert !part.correctTag.equals(token.currentTag);

                // add to list of corrected tokens
                if (changedIndexes != null)
                        localChangedIndexes.add(index);
            }

            if (token.rightTag != null)
            {
                // check if applying the rule really corrects the tag
                if (token.rightTag.equals(part.correctTag))
                {
                    assert token.GetStatus() == Token.Status.ERROR;
                    result.right++;
                }
                else if (token.GetStatus() == Token.Status.CORRECT)
                    result.wrong++;
            }
        }

        if (canApply)
        {
            if (changedIndexes != null)
                changedIndexes.addAll(localChangedIndexes);
            result.total = 1;
            return result;
        }

        // application failed -> return null result
        return new TestResult();
    }

    private void ApplyAtIndex(List<Token> tokens, int startIndex)
    {
        for (int i = 0; i < Length(); i++)
        {
            Part part = parts[i];
            int tokenIndex = startIndex + i;
            Token token = tokens.get(tokenIndex);

            // ensure correctness
            assert token.currentTag.equals(part.originalTag);

            // change the tag to correct one
            if (part.correctTag != null)
            {
                assert token.tagSet.contains(part.correctTag);
                token.currentTag = part.correctTag;
            }
        }
    }

    /**
     * Apply rule to to specific index in the sequence.
     *
     * @param tokens         - sequence of tokens
     * @param startIndex     - offset in the sequence where to start
     * @param reallyApply    - whether to modify tags or just simulate
     * @param changedIndexes - collection to which modified tokens' indexes will be added
     * @return result of applying rule
     */
    private TestResult ApplyToIndex(List<Token> tokens, int startIndex, boolean reallyApply,
                                    Collection<Integer> changedIndexes)
    {
        // check bounds
        assert 0 <= startIndex && startIndex + Length() <= tokens.size();

        // test if rule can be applied and calculate result
        TestResult result = new TestResult();
        boolean canApply = true;
        for (int i = 0; i < Length(); i++)
        {
            Part part = parts[i];
            Token token = tokens.get(startIndex + i);

            // current tag doesn't match - can't apply rule
            if (!token.currentTag.equals(part.originalTag))
            {
                canApply = false;
                break;
            }

            // correct tag not present in list of possible tags - rule can't be applied
            if (part.correctTag != null && !token.tagSet.contains(part.correctTag))
            {
                canApply = false;
                break;
            }

            // can't correct the tag
            if (part.correctTag == null || token.rightTag == null)
                continue;

            // check if applying the rule really corrects the tag
            if (token.rightTag.equals(part.correctTag))
            {
                assert token.GetStatus() == Token.Status.ERROR;
                result.right++;
            }
            else if (token.GetStatus() == Token.Status.CORRECT)
                result.wrong++;
        }

        // apply rule if needed and return result
        if (canApply)
        {
            if (reallyApply)
                // actually apply the rule
                for (int i = 0; i < Length(); i++)
                {
                    Part part = parts[i];
                    int tokenIndex = startIndex + i;
                    Token token = tokens.get(tokenIndex);

                    // ensure correctness
                    assert token.currentTag.equals(part.originalTag);

                    // change the tag to correct one
                    if (part.correctTag != null)
                    {
                        assert token.tagSet.contains(part.correctTag);
                        token.currentTag = part.correctTag;

                        // add to list of corrected tokens
                        if (changedIndexes != null)
                            changedIndexes.add(tokenIndex);
                    }
                }

            result.total = 1;
            return result;
        }

        // application failed -> return null result
        return new TestResult();
    }

    /**
     * Whether to learn partial rules or full rules only.
     */
    public static boolean flagLearnPartialRules = true;

    /**
     * Minimal number of parts in rule.
     */
    public static int minRuleLength = 1;

    /**
     * Maximal number of parts in rule.
     */
    public static int maxRuleLength = 4;

    /**
     * Learn all rules which contain specified token.
     *
     * @param rules      - collection to add new rules into.
     * @param tokens     - sequence of tokens.
     * @param pivotIndex - index of relevant token.
     */
    public static void LearnAllAroundIndex(Collection<Rule> rules, List<Token> tokens,
                                           int pivotIndex)
    {
        // check bounds
        assert 0 <= pivotIndex && pivotIndex < tokens.size();

        /*
        for (int startIndex = Math.max(0, pivotIndex - maxRuleLength + 1); startIndex <= pivotIndex;
             startIndex++)
            for (int endIndex = startIndex + minRuleLength;
                 endIndex <= Math.min(startIndex + maxRuleLength, tokens.size()); endIndex++)
            {
                // ensure correctness of iteration
                assert 0 <= startIndex && endIndex <= tokens.size();

                // check if selected range contains given index
                if (startIndex <= pivotIndex && pivotIndex < endIndex)
                    LearnAllInRange(rules, tokens, startIndex, endIndex);
            }
        */
        for (int startIndex = 0; startIndex <= pivotIndex; startIndex++)
            for (int endIndex = startIndex + minRuleLength; endIndex <= startIndex + maxRuleLength;
                 endIndex++)
            {
                // check if selected range is valid contains given index
                if (0 <= startIndex && startIndex <= pivotIndex
                        && pivotIndex < endIndex && endIndex <= tokens.size())
                    LearnAllInRange(rules, tokens, startIndex, endIndex);
            }
    }

    public static void LearnAll(Collection<Rule> rules, List<Token> tokens)
    {
        // iterate over all possible rule ranges
        for (int startIndex = 0; startIndex < tokens.size(); startIndex++)
            for (int endIndex = startIndex + minRuleLength;
                 endIndex <= Math.min(tokens.size(), startIndex + maxRuleLength); endIndex++)
            {
                // ensure correctness of iteration
                assert 0 <= startIndex && endIndex <= tokens.size();
                LearnAllInRange(rules, tokens, startIndex, endIndex);
            }
    }

    /**
     * Learn all rules from [start, end) range of tokens.
     */
    private static void LearnAllInRange(Collection<Rule> rules, List<Token> tokens, int startIndex,
                                        int endIndex)
    {
        // check bounds
        assert 0 <= startIndex && startIndex < endIndex && endIndex <= tokens.size();

        if (flagLearnPartialRules)
            // learn partial rules for all error tokens in range
            for (int errorIndex = startIndex; errorIndex < endIndex; errorIndex++)
            {
                if (tokens.get(errorIndex).GetStatus() == Token.Status.ERROR)
                {
                    Rule rule = Learn(tokens, startIndex, endIndex, errorIndex);
                    if (rule != null)
                        rules.add(rule);
                }
            }
        else                                                                    
        {
            // learn only one full rule
            Rule rule = Learn(tokens, startIndex, endIndex, -1);
            if (rule != null)
                rules.add(rule);
        }
    }

    /**
     * Learn a rule from [start, end) range of tokens. Partial rule is defined by index of error
     * token which should be corrected.
     *
     * @param errorIndex index of error token to transform in case of partial rule, ignored
     *                   otherwise
     */
    private static Rule Learn(List<Token> tokens, int startIndex, int endIndex,
                              int errorIndex)
    {
        // check bounds
        assert 0 <= startIndex && startIndex < endIndex && endIndex <= tokens.size();

        // for partial rule check if the token has error
        if (flagLearnPartialRules)
        {
            assert startIndex <= errorIndex && errorIndex < endIndex;
            assert tokens.get(errorIndex).GetStatus() == Token.Status.ERROR;
        }

        // build the rule itself
        boolean valid = false;
        int ruleLength = endIndex - startIndex;
        Rule rule = new Rule(ruleLength);
        for (int i = 0; i < ruleLength; i++)
        {
            int tokenIndex = startIndex + i;
            Token token = tokens.get(tokenIndex);
            Part part = rule.parts[i];
            part.originalTag = token.currentTag;
            part.correctTag = null;

            // transform error token
            if (token.GetStatus() == Token.Status.ERROR
                    && (!flagLearnPartialRules || tokenIndex == errorIndex))
            {
                part.correctTag = token.rightTag;
                valid = true;
            }

            // skip 'undefined' token
            if (token.GetStatus() == Token.Status.UNDEFINED)
                // TODO: add flag for invalidation of rule
                continue;
        }
        if (flagLearnPartialRules)
            assert valid;
        if (valid)
            return rule;
        else
            return null;
    }

    /**
     * Write collection of rules into file.
     */
    public static void WriteToFile(Collection<Rule> rules, String filename) throws Throwable
    {
        DocumentBuilderFactory docBuildFactory = DocumentBuilderFactory.newInstance();
        DocumentBuilder docBuilder = docBuildFactory.newDocumentBuilder();
        Document document = docBuilder.newDocument();
        Element rootElement = document.createElement("ruleList");
        document.appendChild(rootElement);
        for (Rule rule : rules)
        {
            Element ruleElement = document.createElement("rule");
            rootElement.appendChild(ruleElement);

            for (Part part : rule.parts)
            {
                Element partElement = document.createElement("part");
                ruleElement.appendChild(partElement);
                Element originalTagElement = part.originalTag.ToXmlElement(document, "originalTag");
                partElement.appendChild(originalTagElement);
                if (part.correctTag != null)
                {
                    Element correctTagElement =
                            part.correctTag.ToXmlElement(document, "correctTag");
                    partElement.appendChild(correctTagElement);
                }
            }
        }

        TransformerFactory transFactory = TransformerFactory.newInstance();
        Transformer transformer = transFactory.newTransformer();
        StreamResult streamResult = new StreamResult(filename);
        DOMSource source = new DOMSource(document);
        transformer.transform(source, streamResult);
    }

    /**
     * Read collection of rules from file.
     */
    public static void ReadFromFile(Collection<Rule> rules, String filename) throws Throwable
    {
        DocumentBuilderFactory factory = DocumentBuilderFactory.newInstance();
        DocumentBuilder builder = factory.newDocumentBuilder();
        Document document = builder.parse(new FileInputStream(filename));

        Element ruleListElement = document.getDocumentElement();
        NodeList ruleNodeList = ruleListElement.getChildNodes();
        for (int i = 0; i < ruleNodeList.getLength(); i++)
        {
            Element ruleElement = (Element) ruleNodeList.item(i);
            NodeList partNodeList = ruleElement.getChildNodes();
            Rule rule = new Rule(partNodeList.getLength());
            for (int j = 0; j < partNodeList.getLength(); j++)
            {
                Element partElement = (Element) partNodeList.item(j);
                NodeList tagNodeList = partElement.getChildNodes();
                rule.parts[j].originalTag = Tag.FromXmlTag((Element) tagNodeList.item(0));
                if (tagNodeList.getLength() > 1)
                    rule.parts[j].correctTag = Tag.FromXmlTag((Element) tagNodeList.item(1));
            }
            rules.add(rule);
        }
    }

    /*
     * Inherited.
     */

    @Override
    public boolean equals(Object object)
    {
        if (!(object instanceof Rule))
            return false;

        Rule rule = (Rule) object;
        return Arrays.equals(this.parts, rule.parts);
    }

    @Override
    public String toString()
    {
        return Utils.join(Arrays.asList(parts), "\n");
    }

    @Override
    public int hashCode()
    {
        return Arrays.hashCode(parts);
    }
}
